// SPDX-License-Identifier: BSD 3-Clause
// Copyright (c) 2024 thinkerhui <thinkerhui@qq.com>
// Copyright (c) 2024 McDj26 <2645370670@qq.com>
// Copyright (c) 2024 qiyf6 <1076753979@qq.com>
// Import necessary modules, you might need to adapt these based on your project structure
use crate::buddy::{buddy_free_pages, buddy_get_pages};
use crate::common::{global_mem, BUDDY_PAGE_SIZE, SLAB_MAX_ORDER};
use crate::slab::{alloc_in_slab, free_in_slab};
use crate::util::{page_to_virt, virt_to_page};

const _SIZE: usize = 1 << SLAB_MAX_ORDER;

fn size_to_page_order(size: usize) -> usize {
    let pg_num = (size + BUDDY_PAGE_SIZE - 1) / BUDDY_PAGE_SIZE;
    let mut order = 0;
    let mut tmp = pg_num;

    while tmp > 1 {
        tmp >>= 1;
        order += 1;
    }

    if pg_num > (1 << order) {
        order += 1;
    }

    order
}

unsafe fn kmalloc(size: usize) -> *mut u8 {
    let order;
    let p_page;

    if size <= _SIZE {
        return alloc_in_slab(size);
    }

    if size <= BUDDY_PAGE_SIZE {
        order = 0;
    } else {
        order = size_to_page_order(size);
    }

    p_page = buddy_get_pages(&global_mem, order);
    page_to_virt(&global_mem, p_page)
}

unsafe fn kzalloc(size: usize) -> *mut u8 {
    let ptr = kmalloc(size);

    // Lack of memory
    if ptr.is_null() {
        return ptr;
    }

    // Zero out the allocated memory
    core::ptr::write_bytes(ptr, 0, size);
    ptr
}

unsafe fn kfree(ptr: *mut u8) {
    let p_page = virt_to_page(&global_mem, ptr);

    if let Some(page) = p_page.as_mut() {
        if let Some(slab) = page.slab {
            free_in_slab(ptr);
            return;
        }
    }

    buddy_free_pages(&global_mem, p_page);
}

unsafe fn get_pages(order: usize) -> *mut u8 {
    let p_page = buddy_get_pages(&global_mem, order);
    page_to_virt(&global_mem, p_page)
}

unsafe fn free_pages(addr: *mut u8) {
    let p_page = virt_to_page(&global_mem, addr);
    buddy_free_pages(&global_mem, p_page);
}
